import torch

class VideoPingPongLoop:
    """
    Create ping-pong loop from input frames (ABCD → ABCDDCBA...)
    Author: cherokee
    """
    
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "images": ("IMAGE",),
                "frame_count": ("INT", {
                    "default": 10, 
                    "min": 1, 
                    "max": 99999,
                    "step": 1
                }),
            },
        }

    RETURN_TYPES = ("IMAGE",)
    RETURN_NAMES = ("images",)
    FUNCTION = "pingpong"
    CATEGORY = "video"

    def sequence_generator(self, num_images, frame_count):
        """生成器函数，按需生成乒乓序列索引"""
        if num_images <= 1:
            for _ in range(frame_count):
                yield 0
            return

        # 计算一个完整的乒乓循环的长度（正向+反向，不重复末尾帧）
        cycle_length = (num_images * 2) - 2

        for i in range(frame_count):
            # 计算在当前循环中的位置
            pos = i % cycle_length
            # 如果在前半段，正向遍历
            if pos < num_images:
                yield pos
            # 如果在后半段，反向遍历
            else:
                yield cycle_length - pos

    def pingpong(self, images, frame_count):
        if len(images) == 0:
            raise ValueError("Input frames cannot be empty!")
        
        if len(images) == 1:
            return (images.repeat(frame_count, 1, 1, 1),)

        # 计算每个批次的大小，根据可用内存动态调整
        available_memory = 1 * 1024 * 1024 * 1024  # 假设有1GB可用内存
        single_frame_size = images[0].element_size() * images[0].nelement()
        batch_size = min(max(1, available_memory // (single_frame_size * 2)), 32)  # 确保至少处理1帧，最多32帧
        
        # 使用生成器直接生成帧
        result_frames = []
        sequence_gen = self.sequence_generator(len(images), frame_count)
        
        try:
            current_batch = []
            for idx in sequence_gen:
                current_batch.append(images[idx])
                
                if len(current_batch) >= batch_size:
                    # 处理当前批次
                    batch_tensor = torch.stack(current_batch)
                    result_frames.append(batch_tensor)
                    current_batch = []
                    
                    # 清理GPU内存
                    if torch.cuda.is_available():
                        torch.cuda.empty_cache()
            
            # 处理剩余的帧
            if current_batch:
                batch_tensor = torch.stack(current_batch)
                result_frames.append(batch_tensor)
            
            # 合并所有批次
            if len(result_frames) == 1:
                final_frames = result_frames[0]
            else:
                final_frames = torch.cat(result_frames, dim=0)
            
            return (final_frames,)
            
        finally:
            # 清理中间变量
            del result_frames
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
